"""Run some games in some scenarios locally.

This is the correct entry to start the program.
To test your agent, the major modifications you need to make are:
1. import your own agent class (you definitely know where it is)
2. instantiate your agent and its opponent (replace the Agent below)

The code looks like:

from ai.agent import Agent

red = Agent()
blue = DemoAgent()
"""
import json
import pickle
import time
import zipfile
import os
import numpy

from ai import Agent
from Layer import utils
from train_env import TrainEnv
from Layer.enc_msg_crack_layer import EncMsgCrackLayer
from Layer.enc_msg_transfer_layer import EncMsgTransferLayer
from Layer.network_attack_layer import NetworkAttackLayer
from Layer.network_defend_layer import NetworkDefendLayer
from Layer.third_assist_layer import ThirdAssistLayer
from info_war_env import InfoWarEnv

RED, BLUE, GREEN = 0, 1, -1


def main():
    run_in_single_agent_mode()
    # run_in_multi_agents_mode()


def run_in_single_agent_mode():
    """
    run demo in single agent mode
    """
    print("running in single agent mode...")
    # instantiate agents and env
    red1 = Agent()
    blue1 = Agent()
    env1 = TrainEnv()
    begin = time.time()

    # get data ready, data can from files, web, or any other sources
    with open("Layer/for_infowar_show.json", encoding='utf8') as f:
        scenario_data = json.load(f)
    with open("data/maps/map_19/basic.json", encoding='utf8') as f:
        basic_data = json.load(f)
    with open('data/maps/map_19/cost.pickle', 'rb') as file:
        cost_data = pickle.load(file)
    see_data = numpy.load("data/maps/map_19/19see.npz")['data']

    # varialbe to build replay
    all_states = []

    # player setup info
    player_info = [{
        "seat": 1,
        "faction": 0,
        "role": 1,
        "user_name": "demo",
        "user_id": 0
    },
        {
            "seat": 11,
            "faction": 1,
            "role": 1,
            "user_name": "demo",
            "user_id": 0
        }]

    # env setup info
    env_step_info = {
        "scenario_data": scenario_data,
        "basic_data": basic_data,
        "cost_data": cost_data,
        "see_data": see_data,
        "player_info": player_info
    }

    # setup env
    state = env1.setup(env_step_info)
    all_states.append(state[GREEN])
    print("Environment is ready.")

    # setup AIs
    red1.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 1,
            "faction": 0,
            "role": 0,
            "user_name": "demo",
            "user_id": 0,
            "state": state,
        }
    )
    blue1.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 11,
            "faction": 1,
            "role": 0,
            "user_name": "demo",
            "user_id": 0,
            "state": state,
        }
    )
    print("agents are ready.")

    # 初始化全局变量
    init_global_var(state)

    # 定义信息战裁决系统
    info_war_env = InfoWarEnv()

    # 给红蓝双方添加layer
    red_enc_msg_crack_layer = EncMsgCrackLayer()
    red_network_attack_layer = NetworkAttackLayer()
    red_enc_msg_transfer_layer = EncMsgTransferLayer()
    red_network_defend_layer = NetworkDefendLayer()
    red_enc_msg_crack_layer.setup(RED)
    red_network_attack_layer.setup(RED)
    red_enc_msg_transfer_layer.setup(RED)
    red_network_defend_layer.setup(RED)

    blue_enc_msg_crack_layer = EncMsgCrackLayer()
    blue_network_attack_layer = NetworkAttackLayer()
    blue_enc_msg_transfer_layer = EncMsgTransferLayer()
    blue_network_defend_layer = NetworkDefendLayer()
    blue_enc_msg_crack_layer.setup(BLUE)
    blue_network_attack_layer.setup(BLUE)
    blue_enc_msg_transfer_layer.setup(BLUE)
    blue_network_defend_layer.setup(BLUE)

    # 定义本轮记录信息战对局信息的日志
    all_logs = []

    # loop until the end of game
    print("steping")
    done = False
    while not done:
        info_actions = []
        info_actions += red_enc_msg_crack_layer.step()
        info_actions += blue_enc_msg_crack_layer.step()
        info_actions += red_network_attack_layer.step()
        info_actions += blue_network_attack_layer.step()
        info_actions += red_enc_msg_transfer_layer.step()
        info_actions += blue_enc_msg_transfer_layer.step()
        info_actions += red_network_defend_layer.step()
        info_actions += blue_network_defend_layer.step()

        logs, msg_actions = info_war_env.step(info_actions)  # logs用于自建web界面的复盘展示；msg_action用于在庙算平台中输出信息战的全局播报消息
        temp = {"time": state[RED]["time"]["cur_step"] + 1}
        logs = {**temp, **logs}  # 这是为了将time作为字典中的第一项
        all_logs.append(logs)

        # 原本的火力战元素。TODO 作战agent中需要根据info_system的情况修改其可获取的信息。

        # TODO 如果被信息窃取：就让对面所有算子都能看到被信息窃取的算子
        # 1、分别遍历红蓝双方算子状态(信息战数据结构)，查出被信息窃取的算子
        ops_status_red = utils.get_ops_status(RED)
        ops_status_blue = utils.get_ops_status(BLUE)
        for op_id, op_status in ops_status_red.items():
            if op_status["has_stealed_infos"] == True:
                # 2、如果被信息窃取了，就给对面所有算子的see_enemy_bop_ids中加上该算子id
                for peer_operator in state[BLUE]["operators"]:
                    peer_operator["see_enemy_bop_ids"].append(op_id)
        for op_id, op_status in ops_status_blue.items():  # 同上
            if op_status["has_stealed_infos"] == True:
                for peer_operator in state[RED]["operators"]:
                    peer_operator["see_enemy_bop_ids"].append(op_id)

        actions = []
        actions += red1.step(state[RED])
        actions += blue1.step(state[BLUE])
        actions.extend(msg_actions)  # 将播报信息战消息的动作添加到动作列表
        state, done = env1.step(actions)

        # update_global_var(state)# 由于聚合/结局，上车/下车动作的存在，state中可见的算子会发送动态变化！

        # TODO 如果有算子死亡了，就将死亡状态同步到信息战的自定义结构体中
        # 1、分别遍历红蓝双方算子(信息战数据结构)，查看它的id是否还存在于operators和passengers中（如果算子死亡了，就会被庙算平台的env踢出operators和passengers）
        # if state[RED]["time"]["cur_step"] == 1500:# 调试用
        #     pass
        ops_status_red = utils.get_ops_status(RED)
        ops_status_blue = utils.get_ops_status(BLUE)
        for op_id, op_status in ops_status_red.items():
            op_status["live_status"] = False  # 先假设已死亡
            for item in state[RED]["operators"] + state[RED]["passengers"]:
                if item["obj_id"] == op_id:
                    op_status["live_status"] = True  # 如果没死亡就同步算子的生命状态并跳出循环，接着判判断下一个算子
                    break
        for op_id, op_status in ops_status_blue.items():  # 同上
            op_status["live_status"] = False
            for item in state[BLUE]["operators"] + state[BLUE]["passengers"]:
                if item["obj_id"] == op_id:
                    op_status["live_status"] = True
                    break
        utils.set_ops_status(RED, ops_status_red)
        utils.set_ops_status(BLUE, ops_status_blue)

        all_states.append(state[GREEN])

    env1.reset()
    red1.reset()
    blue1.reset()

    # 资源释放
    red_enc_msg_crack_layer.reset()
    blue_enc_msg_crack_layer.reset()
    utils.reset()

    print(f"Total time: {time.time() - begin:.3f}s")

    # save replay
    zip_name = f"logs/replays/replay_{begin}.zip"
    if not os.path.exists("logs/replays/"):
        os.makedirs("logs/replays/")
    with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED, compresslevel=9) as z:
        for i, ob in enumerate(all_states):
            data = json.dumps(ob, ensure_ascii=False, separators=(",", ":"))
            z.writestr(f"{begin}/{i}", data)

    # 将日志对象转为为json格式存放到文件中
    json_file_name = f"logs/info/events_{begin}.json"
    if not os.path.exists("logs/info/"):  # 如果文件夹不存在则创建
        os.makedirs("logs/info/")
    with open(json_file_name, 'w') as json_file:
        json.dump(all_logs, json_file, ensure_ascii=False)


def to_json_string(o, indent=None):
    return json.dumps(o, ensure_ascii=False, indent=indent, separators=(',', ':'))


def run_in_multi_agents_mode():
    """
    run demo in multi agent mode
    """
    print("running in multi agent mode...")
    # instantiate agents and env
    red1 = Agent()
    red2 = Agent()
    red3 = Agent()
    blue1 = Agent()
    blue2 = Agent()
    blue3 = Agent()
    env1 = TrainEnv()
    begin = time.time()

    # get data ready, data can from files, web, or any other sources
    with open("data/scenarios/333.json", encoding='utf8') as f:
        scenario_data = json.load(f)
    with open("data/maps/333/333_basic.json", encoding='utf8') as f:
        basic_data = json.load(f)
    with open('data/maps/333/333_cost.pickle', 'rb') as file:
        cost_data = pickle.load(file)
    see_data = numpy.load("data/maps/333/333_see.npz")['data']

    # varialbe to build replay
    all_states = []

    # player setup info
    player_info = [{
        "seat": 1,
        "faction": 0,
        "role": 1,
        "user_name": "red1",
        "user_id": 1
    },
        {
            "seat": 2,
            "faction": 0,
            "role": 0,
            "user_name": "red2",
            "user_id": 2
        },
        {
            "seat": 3,
            "faction": 0,
            "role": 0,
            "user_name": "red3",
            "user_id": 3
        },
        {
            "seat": 11,
            "faction": 1,
            "role": 1,
            "user_name": "blue1",
            "user_id": 11
        },
        {
            "seat": 12,
            "faction": 1,
            "role": 0,
            "user_name": "blue2",
            "user_id": 12
        },
        {
            "seat": 13,
            "faction": 1,
            "role": 0,
            "user_name": "blue3",
            "user_id": 13
        }]

    # env setup info
    env_step_info = {
        "scenario_data": scenario_data,
        "basic_data": basic_data,
        "cost_data": cost_data,
        "see_data": see_data,
        "player_info": player_info
    }

    # setup env
    state = env1.setup(env_step_info)
    all_states.append(state[GREEN])
    print("Environment is ready.")

    # setup AIs
    red1.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 1,
            "faction": 0,
            "role": 1,
            "user_name": "red1",
            "user_id": 1,
            "state": state,
        }
    )
    red2.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 2,
            "faction": 0,
            "role": 0,
            "user_name": "red2",
            "user_id": 2,
            "state": state,
        }
    )
    red3.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 3,
            "faction": 0,
            "role": 0,
            "user_name": "red3",
            "user_id": 3,
            "state": state,
        }
    )
    blue1.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 11,
            "faction": 1,
            "role": 1,
            "user_name": "blue1",
            "user_id": 11,
            "state": state,
        }
    )
    blue2.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 12,
            "faction": 1,
            "role": 0,
            "user_name": "blue2",
            "user_id": 12,
            "state": state,
        }
    )
    blue3.setup(
        {
            "scenario": scenario_data,
            "basic_data": basic_data,
            "cost_data": cost_data,
            "see_data": see_data,
            "seat": 13,
            "faction": 1,
            "role": 0,
            "user_name": "blue3",
            "user_id": 13,
            "state": state,
        }
    )
    print("agents are ready.")

    # loop until the end of game
    print("steping")
    done = False
    while not done:
        actions = []
        actions += red1.step(state[RED])
        actions += red2.step(state[RED])
        actions += red3.step(state[RED])
        actions += blue1.step(state[BLUE])
        actions += blue2.step(state[BLUE])
        actions += blue3.step(state[BLUE])
        state, done = env1.step(actions)
        all_states.append(state[GREEN])

    env1.reset()
    red1.reset()
    red2.reset()
    red3.reset()
    blue1.reset()
    blue2.reset()
    blue3.reset()

    print(f"Total time: {time.time() - begin:.3f}s")

    # save replay
    zip_name = f"logs/replays/replay_{begin}.zip"
    if not os.path.exists("logs/replays/"):
        os.makedirs("logs/replays/")
    with zipfile.ZipFile(zip_name, 'w', zipfile.ZIP_DEFLATED, compresslevel=9) as z:
        for i, ob in enumerate(all_states):
            data = json.dumps(ob, ensure_ascii=False, separators=(",", ":"))
            z.writestr(f"{begin}/{i}", data)


def init_global_var(state):
    """
    用于每轮游戏前初始化全局变量
    """
    network_status_red = {}
    network_status_blue = {}
    ops_status_red = {}
    ops_status_blue = {}
    # network_status_green = {} # TODO 暂时不考虑第三方
    for operator in state[RED]["operators"]:
        network_status_red[operator["obj_id"]] = 0.0
        ops_status_red[operator["obj_id"]] = {
            "live_status": True,  # 生命状态
            "has_disguised_access": False,  # 是否被伪装接入
            "has_implanted_virus": False,  # 是否植入病毒
            "has_stealed_infos": False,  # 是否被信息窃取
            "has_cracked_key": False  # 是否被破解密钥
        }
    for operator in state[RED]["passengers"]:  # 在车上的算子也要进行信息战所需结构体的创建
        network_status_red[operator["obj_id"]] = 0.0
        ops_status_red[operator["obj_id"]] = {
            "live_status": True,  # 生命状态
            "has_disguised_access": False,  # 是否被伪装接入
            "has_implanted_virus": False,  # 是否植入病毒
            "has_stealed_infos": False,  # 是否被信息窃取
            "has_cracked_key": False  # 是否被破解密钥
        }
    utils.set_network_status(RED, network_status_red)
    utils.set_ops_status(RED, ops_status_red)
    for operator in state[BLUE]["operators"]:
        network_status_blue[operator["obj_id"]] = 0.0
        ops_status_blue[operator["obj_id"]] = {
            "live_status": True,  # 生命状态
            "has_disguised_access": False,  # 是否被伪装接入
            "has_implanted_virus": False,  # 是否植入病毒
            "has_stealed_infos": False,  # 是否被信息窃取
            "has_cracked_key": False  # 是否被破解密钥
        }
    for operator in state[BLUE]["passengers"]:  # 在车上的算子也要进行信息战所需结构体的创建
        network_status_blue[operator["obj_id"]] = 0.0
        ops_status_blue[operator["obj_id"]] = {
            "live_status": True,  # 生命状态
            "has_disguised_access": False,  # 是否被伪装接入
            "has_implanted_virus": False,  # 是否植入病毒
            "has_stealed_infos": False,  # 是否被信息窃取
            "has_cracked_key": False  # 是否被破解密钥
        }
    utils.set_network_status(BLUE, network_status_blue)
    utils.set_ops_status(BLUE, ops_status_blue)
    utils.set_current_computing_resources(RED, 200)
    utils.set_current_computing_resources(BLUE, 200)
    utils.set_current_attack_num({
        "red": 0,
        "blue": 0,
    })
    # 将攻击次数、攻击成功次数和通信质量做初始化
    utils.set_current_attack_num({
        "red": 0,
        "blue": 0,
    })
    utils.set_current_attack_success_num({
        "red": 0,
        "blue": 0,
    })
    utils.set_communication_quality({
        "red": 1,
        "blue": 1,
    })
    # 将上一轮是否有通信中断的算子做初始化
    utils.set_last_step_communication_stop({
        "red": False,
        "blue": False,
    })


# def update_global_var(state):# 聚合和解聚动作导致的现有所有算子id出现变更时，同步所有变更到信息战数据体中


if __name__ == "__main__":
    main()
